import matplotlib.pyplot as plt
import numpy as np
import os
import tensorflow as tf

from tensorflow.keras.preprocessing import image_dataset_from_directory

# ## 数据预处理

# ### 数据下载

# 在本教程中，您将使用包含数千个猫和狗图像的数据集。下载并解压缩包含图像的 zip 文件，然后使用 `tf.keras.preprocessing.image_dataset_from_directory` 效用函数创建一个 `tf.data.Dataset` 进行训练和验证。您可以在此[教程](https://tensorflow.google.cn/tutorials/load_data/images)中详细了解如何加载图像。

# In[ ]:


_URL = 'https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip'
path_to_zip = tf.keras.utils.get_file('cats_and_dogs.zip', origin=_URL, extract=True)
print(path_to_zip)
PATH = os.path.join(os.path.dirname(path_to_zip), 'cats_and_dogs_filtered')
print(PATH)
#
train_dir = os.path.join(PATH, 'train')
print(train_dir)
validation_dir = os.path.join(PATH, 'validation')
print(validation_dir)
#
BATCH_SIZE = 32
IMG_SIZE = (160, 160)

train_dataset = image_dataset_from_directory(train_dir,
                                             shuffle=True,
                                             batch_size=BATCH_SIZE,
                                             image_size=IMG_SIZE)
print(train_dataset)
#
# # In[ ]:
#
#
validation_dataset = image_dataset_from_directory(validation_dir,
                                                  shuffle=True,
                                                  batch_size=BATCH_SIZE,
                                                  image_size=IMG_SIZE)

# # 显示训练集中的前九个图像和标签：
#
# # In[ ]:
# class_names = train_dataset.class_names
#
# plt.figure(figsize=(10, 10))
# for images, labels in train_dataset.take(1):
#     for i in range(9):
#         ax = plt.subplot(3, 3, i + 1)
#         plt.imshow(images[i].numpy().astype("uint8"))
#         plt.title(class_names[labels[i]])
#         plt.axis("off")
#
# # 由于原始数据集不包含测试集，因此您需要创建一个。为此，请使用 `tf.data.experimental.cardinality` 确定验证集中有多少批次的数据，然后将其中的 20% 移至测试集。
#
# # In[ ]:
#
#
val_batches = tf.data.experimental.cardinality(validation_dataset)
test_dataset = validation_dataset.take(val_batches // 5)
validation_dataset = validation_dataset.skip(val_batches // 5)
print()
#
# # In[ ]:
#
#
# print('Number of validation batches: %d' % tf.data.experimental.cardinality(validation_dataset))
# print('Number of test batches: %d' % tf.data.experimental.cardinality(test_dataset))
#
# # ### 配置数据集以提高性能
#
# # 使用缓冲预提取从磁盘加载图像，以免造成 I/O 阻塞。要详细了解这种方式，请参阅[数据性能](https://tensorflow.google.cn/guide/data_performance)指南。
#
# # In[ ]:
#
#
# AUTOTUNE = tf.data.AUTOTUNE
#
# train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
# validation_dataset = validation_dataset.prefetch(buffer_size=AUTOTUNE)
# test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)
#
# # ### 使用数据扩充
#
# # 当您没有较大的图像数据集时，最好将随机但现实的转换应用于训练图像（例如旋转或水平翻转）来人为引入样本多样性。这有助于使模型暴露于训练数据的不同方面并减少[过拟合](https://tensorflow.google.cn/tutorials/keras/overfit_and_underfit)。您可以在此[教程](https://tensorflow.google.cn/tutorials/images/data_augmentation)中详细了解数据扩充。
#
# # In[ ]:
#
#
# data_augmentation = tf.keras.Sequential([
#     tf.keras.layers.experimental.preprocessing.RandomFlip('horizontal'),
#     tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
# ])
#
# # 注：当您调用 `model.fit` 时，这些层仅在训练过程中才会处于有效状态。在 `model.evaulate` 或 `model.fit` 中的推断模式下使用模型时，它们处于停用状态。
#
# # 我们将这些层重复应用于同一个图像，然后查看结果。
#
# # In[ ]:
#
#
# for image, _ in train_dataset.take(1):
#     plt.figure(figsize=(10, 10))
#     first_image = image[0]
#     for i in range(9):
#         ax = plt.subplot(3, 3, i + 1)
#         augmented_image = data_augmentation(tf.expand_dims(first_image, 0))
#         plt.imshow(augmented_image[0] / 255)
#         plt.axis('off')
# plt.show()
#
# # ### 重新缩放像素值
# #
# # 稍后，您将下载 `tf.keras.applications.MobileNetV2` 作为基础模型。此模型期望像素值处于 `[-1, 1]` 范围内，但此时，图像中的像素值处于 `[0, 255]` 范围内。要重新缩放这些像素值，请使用模型随附的预处理方法。
#
# # In[ ]:
#
#
# preprocess_input = tf.keras.applications.mobilenet_v2.preprocess_input
#
# # 注：另外，您也可以使用 [Rescaling](https://tensorflow.google.cn/api_docs/python/tf/keras/layers/experimental/preprocessing/Rescaling) 层将像素值从 `[0,255]` 重新缩放为 `[-1, 1]`。
#
# # In[ ]:
#
#
# rescale = tf.keras.layers.experimental.preprocessing.Rescaling(1. / 127.5, offset=-1)
#
# # 注：如果使用其他 `tf.keras.applications`，请确保查阅 API 文档以确定它们是否期望 `[-1,1]` 或 `[0,1]` 范围内的像素，或者使用随附的 `preprocess_input` 函数。
#
# # ## 从预训练卷积网络创建基础模型
# #
# # 您将根据 Google 开发的 **MobileNet V2** 模型来创建基础模型。此模型已基于 ImageNet 数据集进行预训练，ImageNet 数据集是一个包含 140 万个图像和 1000 个类的大型数据集。ImageNet 是一个研究训练数据集，具有各种各样的类别，例如 `jackfruit` 和 `syringe`。此知识库将帮助我们对特定数据集中的猫和狗进行分类。
# #
# # 首先，您需要选择将 MobileNet V2 的哪一层用于特征提取。最后的分类层（在“顶部”，因为大多数机器学习模型的图表是从下到上的）不是很有用。相反，您将按照常见做法依赖于展平操作之前的最后一层。此层被称为“瓶颈层”。与最后一层/顶层相比，瓶颈层的特征保留了更多的通用性。
# #
# # 首先，实例化一个已预加载基于 ImageNet 训练的权重的 MobileNet V2 模型。通过指定 **include_top=False** 参数，可以加载不包括顶部分类层的网络，这对于特征提取十分理想。
#
# # In[ ]:
#
#
# # Create the base model from the pre-trained model MobileNet V2
# IMG_SHAPE = IMG_SIZE + (3,)
# base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
#                                                include_top=False,
#                                                weights='imagenet')
#
# # 此特征提取程序将每个 `160x160x3` 图像转换为 `5x5x1280` 的特征块。我们看看它对一批示例图像做了些什么：
#
# # In[ ]:
#
#
# image_batch, label_batch = next(iter(train_dataset))
# feature_batch = base_model(image_batch)
# print(feature_batch.shape)
#
# # ## 特征提取
# #
# # 在此步骤中，您将冻结在上一步中创建的卷积基，并用作特征提取程序。此外，您还可以在其顶部添加分类器以及训练顶级分类器。
#
# # ### 冻结卷积基
#
# # 在编译和训练模型之前，冻结卷积基至关重要。冻结（通过设置 layer.trainable = False）可避免在训练期间更新给定层中的权重。MobileNet V2 具有许多层，因此将整个模型的 `trainable` 标记设置为 False 会冻结所有这些层。
#
# # In[ ]:
#
#
# base_model.trainable = False
#
# # Let's take a look at the base model architecture
# base_model.summary()
#
# # ### 添加分类头
#
# # 要从特征块生成预测，请使用 `tf.keras.layers.GlobalAveragePooling2D` 层在 `5x5` 空间位置内取平均值，以将特征转换成每个图像一个向量（包含 1280 个元素）。
#
# # In[ ]:
#
#
# global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
# feature_batch_average = global_average_layer(feature_batch)
# print(feature_batch_average.shape)
#
# # 应用 `tf.keras.layers.Dense` 层将这些特征转换成每个图像一个预测。您在此处不需要激活函数，因为此预测将被视为 `logit` 或原始预测值。正数预测 1 类，负数预测 0 类。
#
# # In[ ]:
#
#
# prediction_layer = tf.keras.layers.Dense(1)
# prediction_batch = prediction_layer(feature_batch_average)
# print(prediction_batch.shape)
#
# # 通过使用 [Keras 函数式 API](https://tensorflow.google.cn/guide/keras/functional) 将数据扩充、重新缩放、base_model 和特征提取程序层链接在一起来构建模型。如前面所述，由于我们的模型包含 BatchNormalization 层，因此请使用 training = False。
#
# # In[ ]:
#
#
# inputs = tf.keras.Input(shape=(160, 160, 3))
# x = data_augmentation(inputs)
# x = preprocess_input(x)
# x = base_model(x, training=False)
# x = global_average_layer(x)
# x = tf.keras.layers.Dropout(0.2)(x)
# outputs = prediction_layer(x)
# model = tf.keras.Model(inputs, outputs)
#
# # ### 编译模型
# #
# # 在训练模型前，需要先编译模型。由于存在两个类，并且模型提供线性输出，请将二进制交叉熵损失与 `from_logits=True` 结合使用。
#
# # In[ ]:
#
#
# base_learning_rate = 0.0001
# model.compile(optimizer=tf.keras.optimizers.Adam(lr=base_learning_rate),
#               loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
#               metrics=['accuracy'])
#
# # In[ ]:
#
#
# model.summary()
#
# # MobileNet 中的 250 万个参数被冻结，但在密集层中有 1200 个*可训练*参数。它们分为两个 `tf.Variable` 对象，即权重和偏差。
#
# # In[ ]:
#
#
# len(model.trainable_variables)
